from functools import partial
import os
import argparse
import sys


# dir_path=os.path.dirname(os.path.realpath(__file__))
# sys.path.insert(0,os.path.join(dir_path,'src'))
# sys.path.insert(0,os.path.join(dir_path,'src/codeql_adapter'))
# sys.path.insert(0,os.path.join(dir_path,'src/dataset_adapter'))


import src.codeql_tools.preprocess.token_extractor as token_extractor
import src.codeql_tools.config as config
import src.codeql_tools.codeql_api as codeql_api
import src.codeql_tools.codeql_cmd as codeql_cmd

#from tasks import run_for_varmiuse_model,test,test_train

codeql_init=partial(codeql_api.init,config.index_path,config.query_dir,config.code_source)

def run_for_extract_source_from_dataset():
    token_extractor.default_config()
    token_extractor.main()

def run_for_create_codeql_database():
    
    codeql=config.codeql_path
    db_dir=config.codeql_database_dir
    db_name=config.codeql_database_name
    source=config.code_source
    
    codeql_cmd.create_codeql_db(codeql,db_dir,db_name,source)

def run_for_query_database(param):
    args=['class','func','var','var-write','func-args','var-read','var-subscript']

    targs=['ast-node','ast-edge','basic-block','basic-block-connect','cfg-node','cfg-edge']
    #targs=['basic-block','basic-block-connect','cfg-node']
    uargs=['class','func','var','var-write','func-args','var-read','var-subscript']
    
    if param=="":
        uargs=args
    elif param=="test":
        uargs=targs
    

    db_dir=config.codeql_database_dir
    db_name=config.codeql_database_name
    query_dir=config.query_dir
    codeql=config.codeql_path

    for e in uargs:
        codeql_cmd.query_codeql_db(codeql,os.path.join(db_dir,db_name),e,query_dir)

def run_for_validate(param):
    if param=="":
        codeql_init()
        codeql_api.validate_index()
        codeql_api.validate_token()
    elif param=="only_index":
        codeql_init(True)
        codeql_api.validate_index()
    elif param=="only_token":
        codeql_init()
        codeql_api.validate_token()

def run_for_generate_new_dataset(batch_number=5000):
    codeql_init()
    codeql_api.extract_codeql_sample(config.output_dir,config.output_baseFileName,batch_number)

def main():
    parser=argparse.ArgumentParser(description="program for generate source format from codeql")
    parser.add_argument("-a","---action",type=str,default="process",help='')
    parser.add_argument("-p","--param",type=str,default="",help='')
    parser.add_argument("-batch_number",type=int,default=5000,help="")
    args=parser.parse_args()
    if args.action=="extract_source":
        run_for_extract_source_from_dataset()
    elif args.action=="create_database":
        run_for_create_codeql_database()
    elif args.action=="query_database":
        run_for_query_database(args.param)
    elif args.action=="generate_new_dataset":
        run_for_generate_new_dataset(args.batch_number)
    elif args.action=="process":
        run_for_extract_source_from_dataset()
        run_for_create_codeql_database()
        run_for_query_database(args.param)
        run_for_generate_new_dataset(args.batch_number)
    elif args.action=="init":
        run_for_extract_source_from_dataset()
        run_for_create_codeql_database()
        run_for_query_database(args.param)
    elif args.action=="validate":
        run_for_validate(args.param)
    elif args.action=="train":
        run_for_varmiuse_model()
    elif args.action=="test":
        test()
    elif args.action=="test_train":
        test_train()
if __name__=="__main__":
    main()